import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont


def plot_box(image: np.ndarray, boxes: list, labels: list, colors: list):
    """绘制矩形框"""
    # 标签可能出现中文 使用PIL批量绘制
    img_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)

    # 获取图像尺寸并计算字体缩放比例
    height, width = image.shape[:2]
    max_dimension = max(height, width)
    font_scale = 0.65 if max_dimension <= 2000 else 1.5
    line_width = 4 if max_dimension <= 2000 else 8
    # 设置字体
    font_size = int(30 * font_scale)
    font = ImageFont.truetype("assets/fonts/SourceHanSansCN-Bold.otf", font_size)

    # 批量绘制框和标签
    for i, (box, label, color) in enumerate(zip(boxes, labels, colors)):
        p1 = (int(box[0]), int(box[1]))
        p2 = (int(box[2]), int(box[3]))

        # 绘制矩形框
        draw.rectangle([p1, p2], outline=color, width=line_width)

        # 绘制标签
        bbox = draw.textbbox((0, 0), label, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
        text_x = p1[0] + 3
        text_y = p1[1] - 2 * line_width - text_height
        bg_x1 = text_x - 5
        bg_y1 = text_y - 5
        bg_x2 = text_x + text_width + 5
        bg_y2 = text_y + text_height + 5

        # 绘制背景和文本
        draw.rectangle([(bg_x1, bg_y1), (bg_x2, bg_y2)], fill=color)
        draw.text((text_x, text_y), label, font=font, fill=(255, 255, 255))

    # 转回OpenCV格式
    image = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
    return image


class Plot(object):

    @staticmethod
    def plot_box(image: np.ndarray, boxes: list, labels: list, colors: list):
        """绘制一个矩形框"""

        return plot_box(image, boxes, labels, colors)

    @staticmethod
    def hex_to_rgb(hex_color):
        hex_color = hex_color.lstrip('#')
        return tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4))


PLOT = Plot()